# -*- coding: utf-8 -*-
# @Time    : 2019/12/3 11:05
# @Author  : CHEN Wang
# @File    : stock_price_related.py
# @Comment :
# @Software: PyCharm

"""
脚本说明: 获取股票行情相关数据
"""
import pandas as pd
import math

from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool import time_tool

T_A_STK_DAILY_QUOTE = 'astk_qt_stockdailyquote'  # A股股票不复权行情表
T_A_STK_DAILY_QUOTE_BF = 'astk_qt_stockdailyquote_adjbf'  # A股股票前复权行情表
T_A_STK_DAILY_QUOTE_AF = 'astk_qt_stockdailyquote_adjaf'  # A股股票后复权行情表
T_A_STK_ADJ_FACTOR_BF = 'astk_qt_bffactor'  # A股股票前复权因子表
T_A_STK_ADJ_FACTOR_AF = 'astk_qt_affactor'  # A股股票后复权因子表
T_HK_STK_DAILY_QUOTE = 'hk_qt_stockdailyquote'  # 港股股票不复权行情表
T_HK_STK_DAILY_QUOTE_BF = 'hk_qt_stockdailyquote_adjbf'  # 港股前复权股票行情表
T_HK_STK_DAILY_QUOTE_AF = 'hk_qt_stockdailyquote_adjaf'  # 港股后复权股票行情表
HK_2_A_EXCHANGE_RATE_TABLE = 'astk_bd_hkstexchgrate'  # 港币兑人民币汇率表
HK_2_A_EXCHANGE_DEFAULT_RATE = 0.8338  # 港币兑人民币默认汇率
EARLIEST_START_DATE = '1900-01-01'


def get_hk_2_rmb_exchange_rate(end_date=None):
    """
    得到某一天港币兑人民币汇率，默认最新汇率
    如果表中没有查询到数据，返回默认值0.9126

    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :return: float
    """
    if end_date is None:
        end_date = time_tool.get_today(marker='without_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    conn_stk = db_conn.get_stock_conn()
    # 找到表中离the_day这天最近的交易日期
    date_4_exchange = my_mysql.read_v2(select='max(end_date)', sfrom=HK_2_A_EXCHANGE_RATE_TABLE,
                                       where=[f"end_date <= '{end_date}'", 'exrate_type = 1', 'mkt_type = 1'],
                                       conn=conn_stk).iloc[0, 0]
    # 目前表中最小日期为20161205，如果输入日期小于该日期，返回默认值
    if date_4_exchange is None:
        exchange_rate = HK_2_A_EXCHANGE_DEFAULT_RATE
    else:
        # 市场类型：深港通；汇率类型：结算汇率；买入汇率
        exchange_rate = my_mysql.read_v2(select='buy_exrate', sfrom=HK_2_A_EXCHANGE_RATE_TABLE,
                                         where=[f"end_date = '{date_4_exchange}'", 'exrate_type = 1', 'mkt_type = 1'],
                                         conn=conn_stk).iloc[0, 0]
        # 防止表中数据为null
        if exchange_rate is None or math.isnan(exchange_rate):
            exchange_rate = HK_2_A_EXCHANGE_DEFAULT_RATE
    conn_stk.close()

    return exchange_rate


def get_stock_quote(stock_code=None, start_date=None, end_date=None, only_latest=False,
                    market_type='A', adj_type='n', select=None, convert_to_rmb=False, **kwargs):
    """
    获得股票行情相关数据

    :param str or list or None stock_code: 股票代码，支持单个（str）、多个（list）、所有（None）
                                           如果stock_code非空，则不用传market_type参数，会分别去A股和港股表查询stock_code的数据然后合并
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param str market_type: 市场类型，"A"-A股，"HK"-港股，"both"-A股加港股
    :param str adj_type: 复权类型，"b"-前复权，"a"-后复权，"n"-不复权
    :param list select: 需要返回哪些列，若为None，默认返回所有列
    :param bool convert_to_rmb: 是否要将港股的股票价格转换成人民币，默认不转换，如果要转换，select中有关价格的列请不要重命名
    :param kwargs:
        - conn: 数据库连接
    :return: pd.DataFrame
                - stock_code: 股票代码
                - hqmkt_code: 行情市场代码
                - end_date: 交易日期
                - open: 开盘价
                - high: 最高价
                - low: 最低价
                - close: 收盘价
                - pre_close: 昨收价
                - trade_totalvolumen: 成交总量
                - trade_totalamount: 成交总额
                - updown: 涨跌
                - updown_percent: 涨跌幅
                - ma5_price: 5日均价
                - ma10_price: 10日均价
                - ma20_price: 20日均价
                - ma30_price: 30日均价
                - ma60_price: 60日均价
    """
    conn = kwargs.pop('conn', None)
    tmp_conn = db_conn.get_stock_conn() if conn is None else conn

    if select is None:
        select = ['sec_code', 'trade_date', 'hqmkt_code', 'open', 'high', 'low', 'close',
                  'pre_close', 'trade_totalvolumen', 'trade_totalamount', 'updown', 'updown_percent',
                  'ma5_price', 'ma10_price', 'ma20_price', 'ma30_price', 'ma60_price']
    else:
        select = ['sec_code' if x == 'stock_code' else x for x in select]
        select = ['trade_date' if x == 'end_date' else x for x in select]

    where = []
    if isinstance(stock_code, str) or isinstance(stock_code, list):
        if isinstance(stock_code, str):
            stock_code = [stock_code]
        stock_where_condition = get_lst_condition(stock_code, 'sec_code')
        where.append(stock_where_condition)
    elif stock_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if start_date is None:
        start_date = EARLIEST_START_DATE.replace('-', '')
    else:
        start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    if end_date is None:
        end_date = time_tool.get_today(marker='without_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')
    where.append(f"trade_date >= '{start_date}'")
    where.append(f"trade_date <= '{end_date}'")

    if adj_type == 'n':
        table_suffix = ''
    elif adj_type in ['a', 'b']:
        table_suffix = f'_adj{adj_type}f'
    else:
        LOG.error(f"adj_type目前只支持a、b、n，请重新输入")
        return

    def _inner_func(mkt_type):
        if mkt_type == 'A':
            mkt_type_tmp = 'astk'
        else:
            mkt_type_tmp = 'hk'
        table = f'{mkt_type_tmp}_qt_stockdailyquote' + table_suffix

        if only_latest:
            cols = ', '.join(select)
            sql_str = (f"select {cols} from {table} as a, "
                       f"(select sec_code as code, max(trade_date) as enddate "
                       f"from {table} "
                       f"where trade_date <= '{end_date}' "
                       f"and {stock_where_condition} "
                       f"group by sec_code) as b "
                       f"where a.sec_code = b.code "
                       f"and a.trade_date = b.enddate ")
            df = my_mysql.read_v3(sql_str, tmp_conn)
            # 返回一段时间数据的where语句
        else:
            df = my_mysql.read_v2(select=select, where=where, sfrom=table, conn=tmp_conn)

        if mkt_type == 'HK' and convert_to_rmb:
            convert_list = ['open', 'high', 'low', 'close', 'pre_close', 'ma5_price',
                            'ma10_price', 'ma20_price', 'ma30_price', 'ma60_price']
            convert_list = list(set(select).intersection(set(convert_list)))
            if len(convert_list) > 0:
                exchange_rate = get_hk_2_rmb_exchange_rate(end_date)
                df[convert_list] = df[convert_list] * exchange_rate
        return df

    if market_type in ['A', 'HK'] and stock_code is None:
        res_df = _inner_func(market_type)
    elif market_type == 'both' or stock_code is not None:
        df_1 = _inner_func(mkt_type='A')
        df_2 = _inner_func(mkt_type='HK')
        res_df = pd.concat([df_1, df_2], ignore_index=True)
    else:
        LOG.error(f"market_type目前只支持A、HK、both，请重新输入")
        return

    if conn is None:
        tmp_conn.close()

    if res_df.empty:
        LOG.error(f"没有找到{stock_code}的行情数据")
        return

    if 'sec_code' in res_df.columns:
        res_df = res_df.rename(columns={'sec_code': 'stock_code'})

    if 'trade_date' in res_df.columns:
        res_df = res_df.rename(columns={'trade_date': 'end_date'})
        res_df['end_date'] = res_df['end_date'].apply(time_tool.format_date_str)

    return res_df


def get_stock_adj_factor(stock_code, start_date, end_date, only_latest=False, adj_type='b', select=None):
    """
    获得股票复权因子，目前只支持A股市场

    :param str or list or None stock_code: 股票代码，支持单个（str）、多个（list）、所有（None）
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param str adj_type: 复权类型，"b"-前复权，"a"-后复权
    :param list select: 需要返回哪些列，若为None，默认返回所有列
    :return: pd.DataFrame
                - stock_code: 股票代码
                - hqmkt_code: 行情市场代码
                - end_date: 交易日期
                - ex_div_date: 分红日期 date
                - pre_close_price: 昨收价
                - bf_factor: 前复权因子
                - accum_bf_factor: 累计前复权因子
                - af_factor: 后复权因子
                - accum_af_factor: 累计后复权因子
    """
    if adj_type == 'b':
        table = T_A_STK_ADJ_FACTOR_BF
    elif adj_type == 'a':
        table = T_A_STK_ADJ_FACTOR_AF
    else:
        LOG.error(f"adj_type目前只支持a、b，请重新输入")
        return

    if select is None:
        select = ['sec_code', 'trade_date', f'{adj_type}f_factor', 'hqmkt_code', 'ex_div_date',
                  'pre_close_price', f'{adj_type}f_factor', f'accum_{adj_type}f_factor']
    else:
        select = ['sec_code' if x == 'stock_code' else x for x in select]
        select = ['trade_date' if x == 'end_date' else x for x in select]

    where = []
    if isinstance(stock_code, str) or isinstance(stock_code, list):
        if isinstance(stock_code, str):
            stock_code = [stock_code]
        stock_where_condition = get_lst_condition(stock_code, 'sec_code')
        where.append(stock_where_condition)
    elif stock_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    # 将日期转换成带杠字符串
    if start_date is None:
        start_date = EARLIEST_START_DATE.replace('-', '')
    else:
        start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    if end_date is None:
        end_date = time_tool.get_today(marker='without_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    conn = db_conn.get_stock_conn()

    if only_latest:
        cols = ', '.join(select)
        sql_str = (f"select {cols} from {table} as a, "
                   f"(select sec_code as code, max(trade_date) as enddate "
                   f"from {table} "
                   f"where trade_date <= '{end_date}' "
                   f"and {stock_where_condition} "
                   f"group by sec_code) as b "
                   f"where a.sec_code = b.code "
                   f"and a.trade_date = b.enddate ")
        df = my_mysql.read_v3(sql_str, conn)
        # 返回一段时间数据的where语句
    else:
        where.append(f"trade_date >= '{start_date}'")
        where.append(f"trade_date <= '{end_date}'")
        df = my_mysql.read_v2(select=select, where=where, sfrom=table, conn=conn)

    conn.close()

    if 'sec_code' in df.columns:
        df = df.rename(columns={'sec_code': 'stock_code'})

    if 'trade_date' in df.columns:
        df = df.rename(columns={'trade_date': 'end_date'})
        df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_stock_total_return(start_date, end_date, stock_code=None, **kwargs):
    """
    得到从start_date到end_date的时间区间内，每支股票的区间收益
    如果要获取区间收益，用户一定知道区间起始和结束日期，所以start_date和end_date必须输入
    如果直接取start_date和end_date这两天的收盘价进行计算，那么对于在区间之间新上市的股票，无法计算其区间收益
    目前算法：取这段时间的日行情，向后填充，再向前填充，再用最后一天的收盘价和第一天的收盘价进行计算

    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or list or None stock_code: 需要的股票代码列表
    :return: pd.DataFrame
    """
    conn = kwargs.pop('conn', None)
    stock_hq = get_stock_quote(stock_code=stock_code, start_date=start_date, end_date=end_date,
                               only_latest=False, market_type='both', adj_type='b',
                               select=['stock_code', 'end_date', 'close'], conn=conn)
    if stock_hq is None:
        LOG.error(f"没有找到在{start_date}和{end_date}之间，股票代码为{stock_code}的行情数据，请检查")

    stock_hq = stock_hq.set_index(['stock_code', 'end_date'])['close'].unstack()
    stock_hq = stock_hq.ffill(axis=1)
    stock_hq = stock_hq.bfill(axis=1)
    stock_hq['ret'] = (stock_hq[stock_hq.columns[-1]] /
                       stock_hq[stock_hq.columns[0]] - 1)
    stock_ret = stock_hq['ret'].reset_index()

    return stock_ret


"""
之前的API函数，目前都替换成最新的API函数，以下函数已弃用
"""

# def get_stock_quote(stock_code, start_date, end_date,
#                     market_type='A', adj_type='n', select=None, order_by=None):
#     """
#     获得股票行情相关数据
#
#     :param stock_code: 股票代码，支持单个(str)多个(list)全部(None)
#                        如果stock_code非空，则不用传market_type参数，会分别去A股和港股表查询stock_code的数据然后合并
#     :param str start_date: 开始时间，格式"2020-08-01"，若只查询一天，则start_date=None，该天用end_date传入
#     :param str end_date: 结束时间，格式"2020-08-05"
#     :param str market_type: 市场类型，"A"-A股，"HK"-港股，"both"-A股加港股
#     :param str adj_type: 复权类型，"b"-前复权，"a"-后复权，"n"-不复权
#     :param list select: 需要返回哪些列，若为None，默认返回股票代码、交易日期、收盘价
#     :param order_by: 按哪个字段(str)或哪几个字段(list)升序排列
#     :return: pd.DataFrame
#                 - stock_code: 股票代码
#                 - hqmkt_code: 行情市场代码
#                 - trade_date: 交易日期 str "20200801"
#                 - open: 开盘价
#                 - high: 最高价
#                 - low: 最低价
#                 - close: 收盘价
#                 - pre_close: 昨收价
#                 - average: 均价
#                 - trade_totalvolumen: 成交总量
#                 - trade_totalamount: 成交总额
#                 - updown: 涨跌
#                 - updown_percent: 涨跌幅
#                 - ma5_price: 5日均价
#                 - ma10_price: 10日均价
#                 - ma20_price: 20日均价
#                 - ma30_price: 30日均价
#                 - ma60_price: 60日均价
#     """
#     if select is None:
#         select = ['sec_code', 'trade_date', 'close']
#     else:
#         select = ['sec_code' if x == 'stock_code' else x for x in select]
#
#     where = []
#     if isinstance(stock_code, str):
#         where.append(f'sec_code=\'{stock_code}\'')
#     elif isinstance(stock_code, list):
#         where.append(get_lst_condition(stock_code, 'sec_code'))
#     elif stock_code is None:
#         pass
#     else:
#         LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
#         raise NotImplementedError
#
#     end_date = end_date.replace('-', '')
#     if start_date is None:
#         where.append(f"trade_date = '{end_date}'")
#     else:
#         start_date = start_date.replace('-', '')
#         where.append(f"trade_date >= '{start_date}'")
#         where.append(f"trade_date <= '{end_date}'")
#
#     if adj_type == 'n':
#         table_suffix = ''
#     elif adj_type in ['a', 'b']:
#         table_suffix = f'_adj{adj_type}f'
#     else:
#         LOG.error(f"adj_type目前只支持a、b、n，请重新输入")
#         return
#
#     conn = db_conn.get_stock_conn()
#
#     if market_type in ['A', 'HK'] and stock_code is None:
#         if market_type == 'A':
#             market_type = 'astk'
#         else:
#             market_type = 'hk'
#         table = f'{market_type}_qt_stockdailyquote' + table_suffix
#
#         df = my_mysql.read_v2(select=select, where=where, sfrom=table, order_by=order_by, conn=conn)
#     elif market_type == 'both' or stock_code is not None:
#         table_1 = 'astk_qt_stockdailyquote' + table_suffix
#         table_2 = 'hk_qt_stockdailyquote' + table_suffix
#         df_1 = my_mysql.read_v2(select=select, where=where, sfrom=table_1, order_by=order_by, conn=conn)
#         df_2 = my_mysql.read_v2(select=select, where=where, sfrom=table_2, order_by=order_by, conn=conn)
#         df = pd.concat([df_1, df_2])
#     else:
#         LOG.error(f"market_type目前只支持A、HK、both，请重新输入")
#         return
#
#     conn.close()
#
#     if 'sec_code' in select:
#         df = df.rename(columns={'sec_code': 'stock_code'})
#
#     return df
#
#
# def get_stock_adj_factor(stock_code, start_date, end_date,
#                          adj_type='b', select=None, order_by=None):
#     """
#     获得股票复权因子，目前只支持A股市场
#
#     :param stock_code: 股票代码，支持单个(str)多个(list)全部(None)
#     :param str start_date: 开始时间，格式"2020-08-01"，若只查询一天，则start_date=None，该天用end_date传入
#     :param str end_date: 结束时间，格式"2020-08-05"
#     :param str adj_type: 复权类型，"b"-前复权，"a"-后复权
#     :param list select: 需要返回哪些列，若为None，默认返回股票代码、交易日期、复权因子
#     :param order_by: 按哪个字段(str)或哪几个字段(list)升序排列
#     :return: pd.DataFrame
#                 - stock_code: 股票代码
#                 - hqmkt_code: 行情市场代码
#                 - trade_date: 交易日期 str "20200801"
#                 - ex_div_date: 分红日期 date
#                 - pre_close_price: 昨收价
#                 - bf_factor: 前复权因子
#                 - accum_bf_factor: 累计前复权因子
#                 - af_factor: 后复权因子
#                 - accum_af_factor: 累计后复权因子
#     """
#     if adj_type == 'b':
#         factor_col = 'bf_factor'
#         table = T_A_STK_ADJ_FACTOR_BF
#     elif adj_type == 'a':
#         factor_col = 'af_factor'
#         table = T_A_STK_ADJ_FACTOR_AF
#     else:
#         LOG.error(f"adj_type目前只支持a、b，请重新输入")
#         return
#
#     if select is None:
#         select = ['sec_code', 'trade_date', factor_col]
#     else:
#         select = ['sec_code' if x == 'stock_code' else x for x in select]
#
#     where = []
#     if isinstance(stock_code, str):
#         where.append(f'sec_code=\'{stock_code}\'')
#     elif isinstance(stock_code, list):
#         where.append(get_lst_condition(stock_code, 'sec_code'))
#     elif stock_code is None:
#         pass
#     else:
#         LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
#         raise NotImplementedError
#
#     end_date = end_date.replace('-', '')
#     if start_date is None:
#         where.append(f"trade_date = '{end_date}'")
#     else:
#         start_date = start_date.replace('-', '')
#         where.append(f"trade_date >= '{start_date}'")
#         where.append(f"trade_date <= '{end_date}'")
#
#     conn = db_conn.get_stock_conn()
#
#     df = my_mysql.read_v2(select=select, where=where, sfrom=table, order_by=order_by, conn=conn)
#
#     conn.close()
#
#     if 'sec_code' in select:
#         df = df.rename(columns={'sec_code': 'stock_code'})
#
#     return df


if __name__ == '__main__':
    aaa = get_stock_quote(stock_code=['000001', '00001'], start_date='19000101', end_date='19000102',
                          only_latest=True, market_type='both', adj_type='a', convert_to_rmb=True)
    # bbb = get_stock_adj_factor(None, '2020-12-24', '2020-12-25', False)
    # ccc = get_stock_total_return('20200630', '20200930')
